package com.example.mcp.config;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.node.ArrayNode;
import com.fasterxml.jackson.databind.node.ObjectNode;
import org.springframework.stereotype.Component;

/**
 * OpenAI API响应处理器
 * 用于修复中转服务与Spring AI之间的兼容性问题
 * 
 * @author Generated
 * @version 1.0.0
 */
@Component
public class OpenAiResponseProcessor {
    
    private final ObjectMapper objectMapper = new ObjectMapper();
    
    /**
     * 修复finish_reason枚举值映射问题
     * 将中转服务的大写值转换为Spring AI期望的小写值
     */
    public String processResponse(String originalResponse) {
        try {
            // 先进行简单的字符串替换，处理常见的Gemini API问题
            String fixedResponse = originalResponse
                    .replaceAll("\"finish_reason\":\\s*\"STOP\"", "\"finish_reason\": \"stop\"")
                    .replaceAll("\"finish_reason\":\\s*\"TOOL_CALLS\"", "\"finish_reason\": \"tool_calls\"")
                    .replaceAll("\"finish_reason\":\\s*\"TOOL_CALL\"", "\"finish_reason\": \"tool_call\"")
                    .replaceAll("\"finish_reason\":\\s*\"CONTENT_FILTER\"", "\"finish_reason\": \"content_filter\"")
                    .replaceAll("\"finish_reason\":\\s*\"LENGTH\"", "\"finish_reason\": \"length\"");
            
            // 验证修复后的JSON是否有效
            JsonNode rootNode = objectMapper.readTree(fixedResponse);
            
            // 进一步处理choices数组中的finish_reason
            if (rootNode.has("choices") && rootNode.get("choices").isArray()) {
                ArrayNode choicesArray = (ArrayNode) rootNode.get("choices");
                
                for (JsonNode choice : choicesArray) {
                    if (choice.has("finish_reason")) {
                        String finishReason = choice.get("finish_reason").asText();
                        String mappedReason = mapFinishReason(finishReason);
                        
                        if (!finishReason.equals(mappedReason)) {
                            ((ObjectNode) choice).put("finish_reason", mappedReason);
                            System.out.println("✅ 修复finish_reason: " + finishReason + " -> " + mappedReason);
                        }
                    }
                }
            }
            
            return objectMapper.writeValueAsString(rootNode);
            
        } catch (Exception e) {
            System.err.println("❌ 处理响应时出错: " + e.getMessage());
            System.err.println("原始响应: " + originalResponse.substring(0, Math.min(200, originalResponse.length())));
            return originalResponse; // 返回原始响应
        }
    }
    
    /**
     * 处理流式响应中的单个chunk
     * 专门用于处理流式响应中的finish_reason问题
     */
    public String processStreamChunk(String chunk) {
        try {
            // 直接进行字符串替换，这是最快的处理方法
            String fixedChunk = chunk
                    .replaceAll("\"finish_reason\":\\s*\"STOP\"", "\"finish_reason\": \"stop\"")
                    .replaceAll("\"finish_reason\":\\s*\"TOOL_CALLS\"", "\"finish_reason\": \"tool_calls\"")
                    .replaceAll("\"finish_reason\":\\s*\"TOOL_CALL\"", "\"finish_reason\": \"tool_call\"")
                    .replaceAll("\"finish_reason\":\\s*\"CONTENT_FILTER\"", "\"finish_reason\": \"content_filter\"")
                    .replaceAll("\"finish_reason\":\\s*\"LENGTH\"", "\"finish_reason\": \"length\"");
            
            // 验证修复后的JSON是否有效
            objectMapper.readTree(fixedChunk);
            
            return fixedChunk;
            
        } catch (Exception e) {
            System.err.println("❌ 处理流式chunk时出错: " + e.getMessage());
            return chunk; // 返回原始chunk
        }
    }
    
    /**
     * 映射finish_reason值
     */
    private String mapFinishReason(String originalReason) {
        if (originalReason == null) {
            return "stop";
        }
        
        return switch (originalReason.toUpperCase()) {
            case "STOP" -> "stop";
            case "MAX_TOKENS" -> "length";
            case "TOOL_CALL" -> "tool_call";
            case "TOOL_CALLS" -> "tool_calls";
            case "CONTENT_FILTER" -> "content_filter";
            case "LENGTH" -> "length";
            default -> "stop"; // 默认值
        };
    }
}
